#include <iostream>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <atomic>
#include "tracking.h"
#include "../detection/detect.h"
#include "../processing/processing.h"

using namespace std;
using namespace cv;

static cv::Mat image;
vector<cv::Rect> bboxes;
vector<Droplet> droplets;
std::mutex mut;
std::condition_variable cvar;
std::atomic<int> finishedNum = 0;

MultipleObjectTracker::~MultipleObjectTracker() {
    this->close();
}

void MultipleObjectTracker::close() {
    if (trackers.empty()) return;
    for (auto& item : trackers) {
        auto& tracker = item.second;
        delete tracker;
    }
    trackers.clear();
    droplets.clear();
    bboxes.clear();
}

bool bboxCompare(cv::Rect& A, cv::Rect& B) {
    int delta = 10;
    int x1 = A.x / delta;
    int y1 = A.y / delta;
    int x2 = B.x / delta;
    int y2 = B.y / delta;
    return x1 <= x2 && y1 <= y2;
}


void task(CFTracker* tracker, int id) {
    while (true) {
        std::unique_lock <std::mutex> lock(mut);
        //printf("Tracker [%d] wait ... \n", id);
        cvar.wait(lock);
        lock.unlock();
        //printf("Tracker [%d] update ... \n", id);
        tracker->update(image);
        bboxes[id] = tracker->getBBox();
        finishedNum++;
        //printf("Tracker [%d] done.\n", id);
    }
}

void MultipleObjectTracker::init(cv::Mat& frame) {
    this->close();
    manuDetect(frame, bboxes);
    std::sort(bboxes.begin(), bboxes.end(), bboxCompare);
    num = (uint)bboxes.size();
    printf("initTracker: \n");
    for (uint i = 0; i < num; ++i) {
        auto& id = i;
        auto& bbox = bboxes[i];
        printf("  %d: ", id);
        cout << bbox << endl;
        droplets.push_back(Droplet(i, bbox));
        this->trackers[id] = new CFTracker;
        this->trackers[id]->use_ascn = true;
        this->trackers[id]->init(frame, bbox);
    }
    vector<thread> threads(num);
    for (uint i = 0; i < num; ++i) {
        threads[i] = std::thread(task, this->trackers[i], i);
    }
    for (auto& th : threads)   th.detach();
}

void MultipleObjectTracker::reinit(cv::Mat& frame) {
    this->init(frame);
}

void MultipleObjectTracker::update(cv::Mat& frame) {
    image = frame;
    finishedNum = 0;
    cvar.notify_all();
    while (finishedNum != num) {
        std::this_thread::sleep_for(std::chrono::milliseconds(2));
    }
    for (uint i = 0; i < num; ++i) {
        droplets[i].bbox = bboxes[droplets[i].id];
        auto& id = droplets[i].id;
        auto& bbox = droplets[i].bbox;
        cv::Rect tbox{ bbox.x - 2,bbox.y - 25,bbox.width + 4,25 };
        cv::rectangle(frame, tbox, GREEN, -1);
        cv::rectangle(frame, bbox, GREEN, 2, 8, 0);
        cv::putText(frame, to_string(id), cv::Point(bbox.x, bbox.y - 2), 0, 0.6, WHITE, 2);
    }
    cv::imshow("tracking", frame);
}
